package qdrant

import (
	"context"

	qdrant_client "github.com/qdrant/go-client/qdrant"
	"github.com/samber/lo"
	"google.golang.org/grpc"
)

// Orm 支持 基础的Qdrant增删改查操作, T 是存取的数据类型
//
// 使用者需要实现 Valuer 和 Scanner
type Orm[T any] struct {
	conn           *grpc.ClientConn
	collectionName string
	scanner        Scanner[T]
	valuer         Valuer[T]
	groupScanner   GroupScanner[T]
}

// Valuer 将数据转换为 qdrant point
type Valuer[T any] func(T) *qdrant_client.PointStruct

// Scanner 将 qdrant point 转换为类型 T
type Scanner[T any] func(p QdrantResult, withPayload, withVector bool) T

type GroupScanner[T any] func(g *qdrant_client.PointGroup, withPayload, withVector bool) ScoredSearchResult[T]

// QdrantResult 从qdrant中查询到的原始数据，比如 RetrievedPoint 和 ScoredPoint
type QdrantResult interface {
	GetId() *qdrant_client.PointId
	GetPayload() map[string]*qdrant_client.Value
	GetVectors() *qdrant_client.Vectors
}

func New[T any](
	conn *grpc.ClientConn,
	collectionName string,
	scanner Scanner[T],
	valuer Valuer[T],
	groupScanner GroupScanner[T],
) *Orm[T] {
	return &Orm[T]{
		conn:           conn,
		collectionName: collectionName,
		scanner:        scanner,
		valuer:         valuer,
		groupScanner:   groupScanner,
	}
}

func (orm *Orm[T]) GetConn() *grpc.ClientConn {
	return orm.conn
}

func (orm *Orm[T]) GetCollectionName() string {
	return orm.collectionName
}

func (orm *Orm[T]) Upsert(ctx context.Context, datas []T, wait bool) error {
	points := lo.Map(datas, func(t T, _ int) *qdrant_client.PointStruct {
		return orm.valuer(t)
	})

	cli := qdrant_client.NewPointsClient(orm.conn)
	req := &qdrant_client.UpsertPoints{
		CollectionName: orm.collectionName,
		Points:         points,
		Wait:           Ptr(wait),
	}
	_, err := cli.Upsert(ctx, req)
	return err
}

func (orm *Orm[T]) Get(
	ctx context.Context, ids []string, withPayload, withVector bool,
) ([]T, error) {
	cli := qdrant_client.NewPointsClient(orm.conn)
	resp, err := cli.Get(ctx, &qdrant_client.GetPoints{
		CollectionName: orm.collectionName,
		Ids: lo.Map(ids, func(id string, _ int) *qdrant_client.PointId {
			return &qdrant_client.PointId{
				PointIdOptions: &qdrant_client.PointId_Uuid{
					Uuid: id,
				},
			}
		}),
		WithPayload: &qdrant_client.WithPayloadSelector{
			SelectorOptions: &qdrant_client.WithPayloadSelector_Enable{
				Enable: withPayload,
			},
		},
		WithVectors: &qdrant_client.WithVectorsSelector{
			SelectorOptions: &qdrant_client.WithVectorsSelector_Enable{
				Enable: withVector,
			},
		},
	})
	if err != nil {
		return nil, err
	}

	models := lo.Map(resp.Result, func(p *qdrant_client.RetrievedPoint, _ int) T {
		return orm.scanner(p, withPayload, withVector)
	})

	return models, nil
}

func (orm *Orm[T]) Delete(ctx context.Context, ids []string) error {
	cli := qdrant_client.NewPointsClient(orm.conn)

	_, err := cli.Delete(ctx, &qdrant_client.DeletePoints{
		CollectionName: orm.collectionName,
		Points: &qdrant_client.PointsSelector{
			PointsSelectorOneOf: &qdrant_client.PointsSelector_Points{
				Points: &qdrant_client.PointsIdsList{
					Ids: lo.Map(ids, func(id string, _ int) *qdrant_client.PointId {
						return &qdrant_client.PointId{
							PointIdOptions: &qdrant_client.PointId_Uuid{
								Uuid: id,
							},
						}
					},
					),
				},
			},
		},
	})

	return err
}

type ScoredSearchResult[T any] struct {
	Data  T
	Score float32
}

type SearchOptions struct {
	threashold *float32
}
type SearchOpt func(*SearchOptions)

func SearchWithThreashold(f float32) SearchOpt {
	return func(so *SearchOptions) {
		so.threashold = &f
	}
}

func (orm *Orm[T]) GroupSearch(ctx context.Context, searchVec []float32, filter *qdrant_client.Filter, groupBy string,
	top uint, withPayload, withVector bool,
	opts ...SearchOpt) ([]ScoredSearchResult[T], error) {
	cli := qdrant_client.NewPointsClient(orm.conn)
	o := new(SearchOptions)
	for _, opt := range opts {
		opt(o)
	}
	searchReq := &qdrant_client.SearchPointGroups{
		CollectionName: orm.collectionName,
		Vector:         searchVec,
		Filter:         filter,
		Limit:          uint32(top),
		WithPayload: &qdrant_client.WithPayloadSelector{
			SelectorOptions: &qdrant_client.WithPayloadSelector_Enable{
				Enable: withPayload,
			},
		},
		WithVectors: &qdrant_client.WithVectorsSelector{
			SelectorOptions: &qdrant_client.WithVectorsSelector_Enable{
				Enable: withVector,
			},
		},
		GroupBy:   groupBy,
		GroupSize: 1,
	}
	if o.threashold != nil {
		searchReq.ScoreThreshold = o.threashold
	}
	resp, err := cli.SearchGroups(ctx, searchReq)
	if err != nil {
		return nil, err
	}

	results := lo.Map(resp.Result.Groups, func(p *qdrant_client.PointGroup, _ int) ScoredSearchResult[T] {
		return orm.groupScanner(p, withPayload, withVector)
	})

	return results, nil

}

func (orm *Orm[T]) Search(
	ctx context.Context, searchVec []float32, filter *qdrant_client.Filter,
	top uint, withPayload, withVector bool,
	opts ...SearchOpt,
) ([]ScoredSearchResult[T], error) {
	cli := qdrant_client.NewPointsClient(orm.conn)

	o := new(SearchOptions)
	for _, opt := range opts {
		opt(o)
	}

	searchReq := &qdrant_client.SearchPoints{
		CollectionName: orm.collectionName,
		Vector:         searchVec,
		Limit:          uint64(top),
		Filter:         filter,
		WithPayload: &qdrant_client.WithPayloadSelector{
			SelectorOptions: &qdrant_client.WithPayloadSelector_Enable{
				Enable: withPayload,
			},
		},
		WithVectors: &qdrant_client.WithVectorsSelector{
			SelectorOptions: &qdrant_client.WithVectorsSelector_Enable{
				Enable: withVector,
			},
		},
	}
	if o.threashold != nil {
		searchReq.ScoreThreshold = o.threashold
	}
	resp, err := cli.Search(ctx, searchReq)
	if err != nil {
		return nil, err
	}

	results := lo.Map(resp.Result, func(p *qdrant_client.ScoredPoint, _ int) ScoredSearchResult[T] {
		return ScoredSearchResult[T]{
			Data:  orm.scanner(p, withPayload, withVector),
			Score: p.GetScore(),
		}
	})

	return results, nil
}

func (orm *Orm[T]) Scroll(
	ctx context.Context,
	filter *qdrant_client.Filter,
	offsetID string,
	limit uint32,
	withPayload,
	withVector bool,
) (results []T, nextPageOffset string, err error) {
	cli := qdrant_client.NewPointsClient(orm.conn)
	req := &qdrant_client.ScrollPoints{
		CollectionName: orm.collectionName,
		Filter:         filter,
		Limit:          Ptr(limit),
		WithPayload: &qdrant_client.WithPayloadSelector{
			SelectorOptions: &qdrant_client.WithPayloadSelector_Enable{
				Enable: withPayload,
			},
		},
		WithVectors: &qdrant_client.WithVectorsSelector{
			SelectorOptions: &qdrant_client.WithVectorsSelector_Enable{
				Enable: withVector,
			},
		},
	}
	if offsetID != "" {
		req.Offset = &qdrant_client.PointId{
			PointIdOptions: &qdrant_client.PointId_Uuid{
				Uuid: offsetID,
			},
		}
	}

	resp, err := cli.Scroll(ctx, req)
	if err != nil {
		return nil, "", err
	}

	results = lo.Map(resp.Result, func(p *qdrant_client.RetrievedPoint, _ int) T {
		return orm.scanner(p, withPayload, withVector)
	})

	return results, resp.GetNextPageOffset().GetUuid(), nil
}

type Point[T any] struct {
	ID      string
	Vector  []float32
	Payload T
	Score   float32
}

func Ptr[T any](t T) *T {
	return &t
}
